import torchvision
from ModelStatistics import *
import torch
from torch import nn
from torchvision.models.resnet import ResNet, Bottleneck


class ModelParallelResNet50(ResNet):
    def __init__(self, device1, device2, num_classes=10):
        super().__init__(Bottleneck, [3, 4, 6, 3], num_classes=num_classes)
        # conv1 => bn1 => relu => maxpool => layer1-layer4 => avgpool => fc
        self.device1, self.device2 = device1, device2
        self.seq1 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,
            self.layer1,
            self.layer2).to(self.device1)

        self.seq2 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool).to(self.device2)

        self.fc.to(self.device2)

    def forward(self, x):
        # model parts（layers） 的一个（卡间）串行，
        x = x.to(self.device1)
        x = self.seq2(self.seq1(x).to(self.device2))
        return self.fc(x.view(x.size(0), -1))


def test_model():
    device1 = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
    device2 = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(device1, device2)

    model = ModelParallelResNet50(device1, device2)
    all_param_size(model)

    x = torch.randn([2, 3, 1024, 1024])
    print(model(x).device)
    print(model(x).shape)

    # all_param_name_and_size(model)
    # print_model(model)


# test_model()
